package com.tsing.cli.system.aspect;

import cn.hutool.extra.servlet.JakartaServletUtil;
import com.tsing.cli.common.annotation.Limit;
import com.tsing.cli.common.exception.CustomException;
import com.tsing.cli.common.exception.ResponseCode;
import jakarta.servlet.http.HttpServletRequest;
import lombok.RequiredArgsConstructor;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.List;
import java.util.Objects;

/**
 * 限流切面类
 *
 * @author TheTsing
 */
@Aspect
@Order(1)
@Component
@RequiredArgsConstructor
public class LimitAspect {

    private final RedisTemplate<String, Object> redisTemplate;

    @Pointcut("@annotation(com.tsing.cli.common.annotation.Limit)")
    public void pointcut() {
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
        String key = switch (limitAnnotation.limitType()) {
            case KEY -> limitAnnotation.key();
            case IP -> JakartaServletUtil.getClientIP(request);
            case KEY_AND_IP -> limitAnnotation.key() + "_" + JakartaServletUtil.getClientIP(request);
            case METHOD -> method.getName();
            case METHOD_AND_IP -> method.getName() + "_" + JakartaServletUtil.getClientIP(request);
        };
        if (Boolean.TRUE.equals(redisTemplate.execute(new DefaultRedisScript<>(buildLuaScript(), Boolean.class), List.of(limitAnnotation.prefix() + key), limitAnnotation.count(), limitAnnotation.period()))) {
            return joinPoint.proceed();
        } else {
            throw new CustomException(ResponseCode.SYSTEM_IS_BUSY);
        }
    }

    /**
     * lua限流脚本
     */
    private String buildLuaScript() {
        return """
                local num = redis.call('incr',KEYS[1])
                if tonumber(num) == 1 then
                redis.call('pexpire',KEYS[1],ARGV[2])
                return true
                elseif tonumber(num) > tonumber(ARGV[1]) then
                return false
                else
                return true
                end
                """;
    }

}
